from abc import ABC, abstractmethod
import random
from io import BytesIO
from typing import Tuple, List, Union
from dataclasses import dataclass

from PIL import Image, ImageDraw, ImageFont


@dataclass
class CaptchaConfig:
    font_path: str | None = None
    font_size: int = 30
    width: int = 100
    height: int = 50
    duration: int = 700  # Only for GIF


class NoiseGenerator:
    @staticmethod
    async def draw_noise_line(draw: ImageDraw.ImageDraw) -> None:
        width, height = draw._image.size
        start_x, start_y = random.randint(0, width - 1), random.randint(0, height - 1)
        end_x, end_y = min(random.randint(0, width // 2), width - 1), min(random.randint(0, height // 2), height - 1)
        color = tuple(random.randint(0, 255) for _ in range(3))
        draw.line([(start_x, start_y), (end_x, end_y)], fill=color)

    @staticmethod
    async def draw_noise_circle(draw: ImageDraw.ImageDraw) -> None:
        width, height = draw._image.size
        radius = random.randint(1, height // 4)
        x, y = random.randint(radius, width - radius), random.randint(radius, height - radius)
        outline_color = tuple(random.randint(0, 255) for _ in range(3))
        draw.ellipse([(x - radius, y - radius), (x + radius, y + radius)], outline=outline_color)

    @staticmethod
    async def draw_noise_dot(draw: ImageDraw.ImageDraw) -> None:
        width, height = draw._image.size
        x, y = random.randint(0, width - 1), random.randint(0, height - 1)
        color = tuple(random.randint(0, 255) for _ in range(3))
        draw.point((x, y), fill=color)

    @staticmethod
    async def add_noise(draw: ImageDraw.ImageDraw, width: int, height: int) -> None:
        for _ in range(6):
            await NoiseGenerator.draw_noise_line(draw)
        for _ in range(int(width * height * 0.15)):
            await NoiseGenerator.draw_noise_dot(draw)
        for _ in range(3):
            await NoiseGenerator.draw_noise_circle(draw)


class CaptchaGenerator(ABC):
    def __init__(self, config: CaptchaConfig) -> None:
        self.config = config
        self.font = (
            ImageFont.truetype(config.font_path, config.font_size)
            if config.font_path
            else ImageFont.load_default(size=config.font_size)
        )

    @abstractmethod
    async def generate(self) -> Tuple[bytes, int]:
        pass

    async def generate_math_expression(self) -> Tuple[str, int]:
        a, b = random.randint(1, 9), random.randint(1, 9)
        a, b = max(a, b), min(a, b)
        op = random.choice(["+", "*"])
        result = a + b if op == "+" else a * b
        return f"{a} {op} {b} =", result


class GIFCaptchaGenerator(CaptchaGenerator):
    async def create_frame(self, expression: str, hide_index: int) -> Image.Image:
        background_color = tuple(random.randint(0, 255) for _ in range(3))
        image = Image.new("RGB", (self.config.width, self.config.height), color=background_color)
        draw = ImageDraw.Draw(image)

        await NoiseGenerator.add_noise(draw, self.config.width, self.config.height)

        parts = expression.split()
        parts[hide_index] = "?"
        text = " ".join(parts)

        bbox = draw.textbbox((0, 0), text, font=self.font, align="center")
        text_width, text_height = bbox[2] - bbox[0], bbox[3] - bbox[1]

        code_color = tuple(255 - c for c in background_color)
        text_top = (self.config.height - text_height) // 2 - 10
        text_left = (self.config.width - text_width) // 2
        draw.text((text_left, text_top), text, font=self.font, fill=code_color)

        return image

    async def generate(self) -> Tuple[bytes, int]:
        expression, result = await self.generate_math_expression()

        frames = [
            await self.create_frame(expression, 0),
            await self.create_frame(expression, 2),
        ]

        bytes_io = BytesIO()
        frames[0].save(
            bytes_io,
            format="GIF",
            append_images=frames[1:],
            save_all=True,
            duration=self.config.duration,
            loop=0,
        )
        bytes_io.seek(0)
        return bytes_io.getvalue(), result


class PNGCaptchaGenerator(CaptchaGenerator):
    async def generate(self) -> Tuple[bytes, int]:
        expression, result = await self.generate_math_expression()

        background_color = tuple(random.randint(0, 255) for _ in range(3))
        image = Image.new("RGB", (self.config.width, self.config.height), color=background_color)
        draw = ImageDraw.Draw(image)

        await NoiseGenerator.add_noise(draw, self.config.width, self.config.height)

        # 随机隐藏一个数字或运算符
        parts = expression.split()
        hide_index = random.choice([0, 2])  # 随机隐藏第一个数字或第二个数字
        parts[hide_index] = "?"
        text = " ".join(parts)

        bbox = draw.textbbox((0, 0), text, font=self.font, align="center")
        text_width, text_height = bbox[2] - bbox[0], bbox[3] - bbox[1]

        code_color = tuple(255 - c for c in background_color)
        text_top = (self.config.height - text_height) // 2
        text_left = (self.config.width - text_width) // 2
        draw.text((text_left, text_top), text, font=self.font, fill=code_color)

        bytes_io = BytesIO()
        image.save(bytes_io, format="PNG")
        bytes_io.seek(0)
        return bytes_io.getvalue(), result


# 使用示例:
"""
async def generate_captcha(captcha_type: str = "gif") -> Tuple[bytes, int]:
    config = CaptchaConfig(
        width=100,
        height=50,
        font_size=25,
        duration=700
    )
    
    generator: CaptchaGenerator
    if captcha_type == "gif":
        generator = GIFCaptchaGenerator(config)
    else:
        generator = PNGCaptchaGenerator(config)
        
    return await generator.generate()
"""
